use std::collections::HashMap;
use std::fmt::Display;
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::sync::{mpsc, broadcast};
use tokio::net::{TcpListener, TcpStream};
use tokio::task::JoinHandle;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use futures::stream::{FuturesUnordered, StreamExt};

const BUFFER_LEN: usize = 4096;

#[derive(Debug)]
pub enum Error {
    IoError(std::io::Error),
    MapError(String),
    TuiError(String),
    JsonError(serde_json::error::Error),
}

impl From<std::io::Error> for Error {
    fn from(e: std::io::Error) -> Error {
        Error::IoError(e)
    }
}

impl From<serde_json::error::Error> for Error {
    fn from(e: serde_json::error::Error) -> Error {
        Error::JsonError(e)
    }
}

pub struct Command {}

impl Default for Command {
    fn default() -> Self {
        Command {}
    }
}

pub struct Mapper {
    lst_addr: String,
    lst_port: u16,
    dst_addr: String,
    dst_port: u16,

    exit_cb: Option<Box<dyn FnOnce(&mut Self, &Result<(), Error>) + Send + 'static>>,
    cmd_rx: mpsc::Receiver<Command>,
    cmd_rx_closed: bool
}

impl Mapper {
    pub fn new(lst_addr: &str, lst_port: u16,
               dst_addr: &str, dst_port: u16)
        -> (Mapper, MapperHandle) {

        let lst_addr = lst_addr.into();
        let dst_addr = dst_addr.into();
        let (cmd_tx, cmd_rx) = mpsc::channel(16);
        (
            Mapper {
                lst_addr, lst_port, dst_addr, dst_port,
                exit_cb: None,
                cmd_rx,
                cmd_rx_closed: false
            },
            MapperHandle {
                cmd_tx
            }
        )
    }

    pub fn on_exit<F>(&mut self, f: F) -> &mut Self
        where
            F: FnOnce(&mut Self, &Result<(), Error>) + Send + 'static
    {
        self.exit_cb = Some(Box::new(f));
        self
    }

    pub async fn run(mut self) -> Result<(), Error> {
        let res = self.map().await;
        // On exit, call up the callback
        if let Some(cb) = self.exit_cb {
            self.exit_cb = None;
            cb(&mut self, &res);
        }
        res
    }

    async fn forward_data(mut read_half: OwnedReadHalf,
                          mut write_half: OwnedWriteHalf,
                          mut rx: broadcast::Receiver<()>) -> Result<(), Error> {
        use tokio::sync::broadcast::error::RecvError;
        // Double buffer, but the data would only exists on one of these two buffers
        let mut rd_buffer: Vec<u8> = Vec::with_capacity(BUFFER_LEN);
        rd_buffer.resize(BUFFER_LEN, 0);
        let mut wr_buffer: Vec<u8> = Vec::new();
        std::assert!(rd_buffer.len() == BUFFER_LEN);
        let mut rx_len: usize = 0;
        let mut tx_len: usize = 0;

        loop {
            tokio::select! {
                v = read_half.read(&mut rd_buffer[..]), if rx_len == 0 => {
                    std::assert!(rd_buffer.len() == BUFFER_LEN);
                    match v {
                        Ok(len) => {
                            if len == 0 {
                                // No more data to read, exit
                                break;
                            }
                            rx_len = len;
                            tx_len = 0;
                            // Swap data to write buffer
                            std::mem::swap(&mut rd_buffer, &mut wr_buffer);
                        },
                        Err(e) => {
                            // Failed to read, just exit
                            Err(e)?
                        }
                    }
                },
                v = write_half.write(&wr_buffer[tx_len..rx_len]), if rx_len != 0 => {
                    std::assert!(wr_buffer.len() == BUFFER_LEN);
                    std::assert!(tx_len < rx_len);
                    match v {
                        Ok(len) => {
                            if len == 0 {
                                // Could not write data any more
                                break;
                            }
                            tx_len += len;
                            std::assert!(tx_len <= rx_len);
                            if tx_len == rx_len {
                                // All data has beean written
                                rx_len = 0;
                                tx_len = 0;
                                // Swap the buffer to read buffer
                                std::mem::swap(&mut rd_buffer, &mut wr_buffer);
                            }
                        },
                        Err(e) => {
                            // Failed to write, just exit
                            Err(e)?
                        }
                    }
                },
                v = rx.recv() => {
                    match v {
                        Ok(()) => {
                            // exit
                            // TODO: Handler commands
                            break;
                        },
                        Err(RecvError::Closed) => {
                            // Mapper::run existed before we exist, this should not
                            // happend!
                            std::panic!("Mapper::run exist without waiting us!");
                        },
                        Err(RecvError::Lagged(_)) => {
                            // Nothing we can do for this type of error
                            // TODO: Warn
                        }
                    }
                }
            }
        }

        Ok(())
    }

    async fn mapper(lst_sk: TcpStream, dst_addr: String, dst_port: u16,
                    rx1: broadcast::Receiver<()>, rx2: broadcast::Receiver<()>) -> Result<(), Error> {
        // TODO
        let dst_sk = TcpStream::connect((dst_addr.as_str(), dst_port)).await?;
        let (lst_rd_half, lst_wr_half) = lst_sk.into_split();
        let (dst_rd_half, dst_wr_half) = dst_sk.into_split();

        let jh_lst2dst = tokio::spawn(Self::forward_data(lst_rd_half, dst_wr_half, rx1));
        let jh_dst2lst = tokio::spawn(Self::forward_data(dst_rd_half, lst_wr_half, rx2));

        // Now wait all these two JoinHandle to exist
        let _ = tokio::join!(jh_lst2dst, jh_dst2lst);
        // TODO: Return error if need
        Ok(())
    }

    async fn map(&mut self) -> Result<(), Error> {
        // TODO
        let mut mappers = FuturesUnordered::<JoinHandle<Result<(), Error>>>::new();
        let (broadcast_tx, _) = broadcast::channel::<()>(32);

        let lst = TcpListener::bind((self.lst_addr.as_str(), self.lst_port)).await?;

        loop {
            tokio::select! {
                r = lst.accept() => {
                    match r {
                        Ok((src_sk, _)) => {
                            mappers.push(
                                tokio::spawn(
                                    Self::mapper(src_sk,
                                                   self.dst_addr.clone(), self.dst_port,
                                                   broadcast_tx.subscribe(),
                                                   broadcast_tx.subscribe())
                            ));
                        },
                        // Failed to accept, exist
                        Err(e) => {
                            // TODO: throw error
                            println!("accept failed {:?}", e);
                            break;
                        }
                    }
                },
                _ = mappers.next(), if !mappers.is_empty() => {
                    // Some mapper has existed, do nothing
                    // TODO: Report if the mapper exist with error
                },
                v = self.cmd_rx.recv(), if !self.cmd_rx_closed => {
                    // We should stop now
                    match v {
                        Some(_) => {
                            println!("recv from cmd_rx");
                            // TODO: Handle commands
                            break;
                        },
                        None => {
                            println!("cmd_rx is closed");
                            self.cmd_rx_closed = true;
                            // Nothing we can do here
                        }
                    }
                }
            }
        }

        // Now, tell everybody that we have stopped
        let _ = broadcast_tx.send(());
        // Wait all other guys exit
        // FIXME: Use collect?
        while !mappers.is_empty() {
            mappers.next().await;
        }
        Ok(())
    }
}

pub struct MapperHandle {
    cmd_tx: mpsc::Sender<Command>
}

impl MapperHandle {
    pub async fn shutdown(self) {
        if self.cmd_tx.is_closed() {
            return;
        }
        // Tell Mapper to exist
        let _ = self.cmd_tx.send(Default::default()).await;
        // Wait Mapper to be exited
        self.cmd_tx.closed().await;
    }

    pub fn is_exited(&self) -> bool {
        self.cmd_tx.is_closed()
    }
}

impl Clone for MapperHandle {
    fn clone(&self) -> Self {
        MapperHandle {
            cmd_tx: self.cmd_tx.clone()
        }
    }
}

#[derive(Clone, Hash, PartialEq, Eq)]
pub struct MapInfo {
    pub lst_addr: String,
    pub lst_port: u16,
    pub dst_addr: String,
    pub dst_port: u16
}

impl MapInfo {
    pub fn new(lst_addr: &str, lst_port: u16, dst_addr: &str, dst_port: u16) -> Self {
        MapInfo {
            lst_addr: lst_addr.into(),
            lst_port,
            dst_addr: dst_addr.into(),
            dst_port,
        }
    }
}

impl Display for MapInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::write!(f, "{}:{} -> {}:{}", self.lst_addr, self.lst_port,
                                       self.dst_addr, self.dst_port)
    }
}

enum Op {
    Map(MapInfo, Box<dyn FnOnce() + Send + 'static>),
    Unmap(MapInfo)
}

struct Msg {
    replier: Option<mpsc::Sender<OpReply>>,
    op: Op
}

impl Msg {
    fn new(replier: Option<mpsc::Sender<OpReply>>, op: Op) -> Self {
        Msg {replier, op}
    }
}

enum OpReply {
    Succeed,
    Fail(String)
}


#[allow(dead_code)]
pub struct MapperRuntime {
    op_tx: mpsc::Sender<Msg>,
    rt_thread: std::thread::JoinHandle<()>
}

impl MapperRuntime {
    pub fn map<F>(&self, info: MapInfo, on_exit: F) -> Result<(), Error>
        where F: FnOnce() + Send + 'static {
        match self.send_op(Op::Map(info, Box::new(on_exit)))? {
            OpReply::Succeed => Ok(()),
            OpReply::Fail(m) => Err(Error::MapError(m))
        }
    }

    pub fn unmap(&self, info: MapInfo) -> Result<(), Error> {
        match self.send_op(Op::Unmap(info))? {
            OpReply::Succeed => Ok(()),
            OpReply::Fail(m) => Err(Error::MapError(m))
        }
    }

    fn send_op(&self, op: Op) -> Result<OpReply, Error> {
        let (tx, mut rx) = mpsc::channel(32);
        match self.op_tx.blocking_send(Msg::new(Some(tx), op)) {
            Ok(_) => {
                if let Some(reply) = rx.blocking_recv() {
                    Ok(reply)
                } else {
                    Err(Error::MapError("No reply from mapper".into()))
                }
            },
            Err(_) => {
                Err(Error::MapError("Failed to send op".into()))
            }
        }
    }
}

pub struct  MapperRuntimeBuilder;

impl MapperRuntimeBuilder {
    pub fn new() -> Self {
        MapperRuntimeBuilder {}
    }

    pub fn build(self) -> MapperRuntime {
        let (op_tx, op_rx) = mpsc::channel(32);
        let rt_thread = std::thread::spawn(move || {
            MapperRuntimeInternal::new().run(op_rx)
        });

        MapperRuntime {op_tx, rt_thread}
    }
}

struct MapperRuntimeInternal {
    mappers: HashMap<MapInfo, MapperHandle>,
}

impl MapperRuntimeInternal {
    fn new() -> Self {
        MapperRuntimeInternal {
            mappers: HashMap::new()
        }
    }

    fn run(mut self, mut op_rx: mpsc::Receiver<Msg>) {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async move {
            loop {
                if let Some(msg) = op_rx.recv().await {
                    if !self.handle_msg(msg).await {
                        break;
                    }
                } else {
                    // For None, we should close the rt
                    // TODO: Exit graceful
                    break;
                }
            }
        });
    }

    // Return: true, continue; false, exist.
    async fn handle_msg(&mut self, msg: Msg) -> bool {
        let Msg {replier, op} = msg;
        match op {
            Op::Map(info, cb) => {
                self.handle_map(replier, info, cb).await
            },
            Op::Unmap(info) => {
                self.handle_unmap(replier, info).await
            }
        }
    }

    // You must unmap if you want to map again even if you failed to map, otherwise
    // Map would return error.
    async fn handle_map(&mut self, replier: Option<mpsc::Sender<OpReply>>,
                        info: MapInfo, cb: Box<dyn FnOnce() + Send + 'static>) -> bool {
        let mut op_reply = OpReply::Succeed;
        if self.mappers.contains_key(&info) {
            op_reply = OpReply::Fail("Map exist".into());
        } else {
            let (mut mapper, mapper_handle) = Mapper::new(&info.lst_addr, info.lst_port,
                                                      &info.dst_addr, info.dst_port);
            // Insert MapHandle first
            self.mappers.insert(info, mapper_handle);
            tokio::spawn(async move {
                mapper.on_exit(move |_, _| {
                    // Just callup the callback
                    cb()
                });
                let _ = mapper.run().await;
            });
        }
        if let Some(replier) = replier {
            // Do nothing even if we failed to send the reply
            let _ = replier.send(op_reply).await;
        }
        true
    }

    async fn handle_unmap(&mut self, replier: Option<mpsc::Sender<OpReply>>,
                          info: MapInfo) -> bool {
        let mut op_reply = OpReply::Succeed;
        if !self.mappers.contains_key(&info) {
            op_reply = OpReply::Fail("Map not exist".into());
        } else {
            let mapper_handle = self.mappers.remove(&info).unwrap();
            mapper_handle.shutdown().await;
        }
        if let Some(replier) = replier {
            let _ = replier.send(op_reply).await;
        }
        true
    }
}
